import numpy as np
from metrics import *
import _pickle
import saliency_metrics

attentionPath = 'D:/VR_project/PanoSaliency/data/'
attentionList = ["saliency_ds2_topic0", "saliency_ds2_topic1", "saliency_ds2_topic2",
                "saliency_ds2_topic3", "saliency_ds2_topic4", "saliency_ds2_topic5",
                "saliency_ds2_topic6", "saliency_ds2_topic7", "saliency_ds2_topic8"]

samplePath = 'D:/VR_project/ViewPrediction/frames/attentionSaliency/'
sampleDict = {0: 'Conan1', 1: 'Skiing', 2: 'Alien', 3: 'Conan2', 4: 'Surfing',
                5: 'War', 6: 'Cooking', 7: 'Football', 8: 'Rhinos'}

# load
n = 0
l = len(attentionList)
CC_list = []
NSS_list = []
KL_list = []
AUC_list = []
sAUC_list = []
for i in range(len(attentionList)):
    saliencyPath = attentionPath + attentionList[i]
    try:
        saliency_array = np.array(pickle.load(open(saliencyPath, 'rb'), encoding='bytes'), dtype=object)
    except _pickle.UnpicklingError:
        saliency_array = np.load(saliencyPath, allow_pickle=True)

    t_list = [item[0] for item in saliency_array]
    fixation_maps = np.array([create_fixation_map(item[1]) for item in saliency_array])
    saliency_maps = np.array([item[2] for item in saliency_array])
    # print(attentionList[i], fixation_maps.shape, saliency_maps.shape)

    CC_list_list, NSS_list_list, KL_list_list, AUC_list_list = [], [], [], []
    sAUC_list_list = []
    for idx in range(0, len(saliency_maps)):
        fixmap = fixation_maps[idx]
        salmap = saliency_maps[idx]
        if len(np.unique(fixmap)) != 2 or fixmap.shape != salmap.shape:
            continue
        CC = cal_CC(fixmap, salmap)
        NSS = cal_NSS(fixmap, salmap)
        KL = cal_KL(fixmap, salmap)
        AUC = cal_AUC(fixmap, salmap)
        sAUC = saliency_metrics.auc_shuff(salmap, fixmap, fixmap)

        if CC == -1 or NSS == -1:
            continue
        if KL == float('inf'):
            continue
        CC_list_list.append(CC)
        NSS_list_list.append(NSS)
        KL_list_list.append(KL)
        AUC_list_list.append(AUC)
        sAUC_list_list.append(sAUC)

    CC_list.append(np.mean(CC_list_list))
    NSS_list.append(np.mean(NSS_list_list))
    KL_list.append(np.mean(KL_list_list))
    AUC_list.append(np.mean(AUC_list_list))
    sAUC_list.append(np.mean(sAUC_list_list))

    print(f'{sampleDict[i]} CC={round(CC_list[i], 4)}, NSS={round(NSS_list[i], 4)}, '
          f'KL={round(KL_list[i], 4)}, AUC={round(AUC_list[i], 4)}, sAUC={round(sAUC_list[i], 4)}')


"""path = 'dataset2.npy'
saliency_array = np.load(attentionPath + path, allow_pickle=True)
print('dataset2.npy', saliency_array.shape)
for i in range(0, len(saliency_array), 5):
    x = saliency_array[i]
    plt.imshow(x)
    plt.show()
"""
print("Done!")